import torch
import numpy as np
# 模型定义（你自己的结构）
import sys
import os
sys.path.append(os.path.dirname(__file__))
from MILmodel import COPDNet


# class MyBinaryClassifier(torch.nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.model = torch.nn.Sequential(
#             torch.nn.Flatten(),
#             torch.nn.Linear(3 * 224 * 224, 128),
#             torch.nn.ReLU(),
#             torch.nn.Linear(128, 2)  # 二分类
#         )
#
#     def forward(self, x):
#         return self.model(x)

# 加载模型权重
model = COPDNet(model_name="resnet34")
model_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "models", "fold_4_model_resnet34_val_acc_1.0000_val_auc_1.0000_train_acc_0.9542_train_auc_0.9888_epoch_5.pth")
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()

# # 推理函数：输入为 numpy 数组（.npy 文件加载的内容）
# def predict(input_array: np.ndarray):
#     # 选择第 0 张图像
#     if input_array.ndim == 4:
#         input_array = input_array[0]  # 变成 (224, 224, 3)
#
#     # 转换为 float32 并归一化
#     input_array = input_array.astype(np.float32) / 255.0
#
#     # HWC → CHW
#     input_array = np.transpose(input_array, (2, 0, 1))
#
#     # 转 tensor + 添加 batch 维度
#     tensor = torch.from_numpy(input_array).unsqueeze(0)
#
#     with torch.no_grad():
#         logits = model(tensor)
#         probs = torch.softmax(logits[0], dim=0)
#
#     pred_class = torch.argmax(probs).item()
#     confidence = float(probs[pred_class])
#
#     return {
#         "label": str(pred_class),  # "0" or "1"
#         "confidence": confidence,
#         "probabilities": {
#             "0": float(probs[0]),
#             "1": float(probs[1])
#         }
#     }
def predict(input_array: np.ndarray):
    import torch

    if input_array.ndim == 4 and input_array.shape[-1] == 3:
        # 多张图像：B, H, W, C → B, C, H, W
        input_array = np.transpose(input_array, (0, 3, 1, 2))
    elif input_array.ndim == 3 and input_array.shape[-1] == 3:
        # 单张图像：H, W, C → 1, C, H, W
        input_array = np.transpose(input_array, (2, 0, 1))
        input_array = np.expand_dims(input_array, 0)
    else:
        raise ValueError(f"输入维度不合法，接收到 shape: {input_array.shape}")

    input_array = input_array.astype(np.float32) / 255.0
    tensor = torch.from_numpy(input_array)

    with torch.no_grad():
        logits = model(tensor)
        probs = torch.softmax(logits[0], dim=0)

    pred_class = torch.argmax(probs).item()
    confidence = float(probs[pred_class])

    return {
        "label": str(pred_class),
        "confidence": confidence,
        "probabilities": {
            "0": float(probs[0]),
            "1": float(probs[1])
        }
    }
